import pickle
import networkx as nx
import os
from collections import defaultdict, Counter


def load_graph(pickle_file):
    """Load graph from pickle file"""
    try:
        with open(pickle_file, "rb") as f:
            G = pickle.load(f)
        return G
    except Exception as e:
        print(f"Error loading file: {e}")
        return None


def analyze_node_relationships(G, max_examples=10):
    """Analyze and print node relationships by type"""
    if not isinstance(G, (nx.Graph, nx.DiGraph, nx.MultiGraph, nx.MultiDiGraph)):
        print("Not a NetworkX graph object")
        return

    print("=== NODE RELATIONSHIP ANALYSIS ===\n")

    # Get all node types
    node_types = {}
    for node in G.nodes():
        node_type = G.nodes[node].get("type", "unknown")
        node_types[node] = node_type

    type_counts = Counter(node_types.values())
    print("Node Type Summary:")
    for node_type, count in sorted(type_counts.items()):
        print(f"  {node_type}: {count} nodes")
    print()

    # Analyze relationships between different node types
    relationship_types = defaultdict(list)

    # For MultiGraph, we need to handle edges with keys
    if isinstance(G, (nx.MultiGraph, nx.MultiDiGraph)):
        edges = G.edges(keys=True, data=True)
        for u, v, key, data in edges:
            u_type = node_types.get(u, "unknown")
            v_type = node_types.get(v, "unknown")
            rel_type = f"{u_type} -> {v_type}"
            relationship_types[rel_type].append((u, v, key, data))
    else:
        edges = G.edges(data=True)
        for u, v, data in edges:
            u_type = node_types.get(u, "unknown")
            v_type = node_types.get(v, "unknown")
            rel_type = f"{u_type} -> {v_type}"
            relationship_types[rel_type].append((u, v, data))

    print("Relationship Type Summary:")
    for rel_type, edges in sorted(
        relationship_types.items(), key=lambda x: len(x[1]), reverse=True
    ):
        print(f"  {rel_type}: {len(edges)} edges")
    print()

    # Show examples for each relationship type
    print("=== DETAILED RELATIONSHIP EXAMPLES ===\n")

    for rel_type, edges in sorted(
        relationship_types.items(), key=lambda x: len(x[1]), reverse=True
    ):
        print(f"--- {rel_type} ({len(edges)} total) ---")

        # Show up to max_examples
        for i, edge_info in enumerate(edges[:max_examples]):
            if isinstance(G, (nx.MultiGraph, nx.MultiDiGraph)):
                u, v, key, data = edge_info
                print(f"  {i+1}. {u} -> {v} (key: {key})")
                if data:
                    print(f"     Edge attributes: {data}")
            else:
                u, v, data = edge_info
                print(f"  {i+1}. {u} -> {v}")
                if data:
                    print(f"     Edge attributes: {data}")

        if len(edges) > max_examples:
            print(f"     ... and {len(edges) - max_examples} more")
        print()


def analyze_specific_node_types(G, target_types=None):
    """Analyze specific node types in detail"""
    if target_types is None:
        target_types = ["file", "function", "class", "directory"]

    print("=== SPECIFIC NODE TYPE ANALYSIS ===\n")

    for node_type in target_types:
        nodes_of_type = [
            node for node in G.nodes() if G.nodes[node].get("type") == node_type
        ]

        if not nodes_of_type:
            continue

        print(f"--- {node_type.upper()} NODES ({len(nodes_of_type)} total) ---")

        # Show some examples
        for i, node in enumerate(nodes_of_type[:3]):
            print(f"  {i+1}. {node}")

            # Show node attributes
            attrs = G.nodes[node]
            if attrs:
                for key, value in attrs.items():
                    if key != "type":
                        print(f"     {key}: {value}")

            # Show connections
            if G.is_directed():
                in_edges = list(G.in_edges(node))
                out_edges = list(G.out_edges(node))
                print(f"     In-degree: {len(in_edges)}, Out-degree: {len(out_edges)}")

                if in_edges:
                    print(f"     Sample incoming: {in_edges[:2]}")
                if out_edges:
                    print(f"     Sample outgoing: {out_edges[:2]}")
            else:
                neighbors = list(G.neighbors(node))
                print(f"     Degree: {len(neighbors)}")
                if neighbors:
                    print(f"     Sample neighbors: {neighbors[:3]}")
            print()

        if len(nodes_of_type) > 10:
            print(f"     ... and {len(nodes_of_type) - 10} more {node_type} nodes")
        print()


def analyze_high_degree_nodes(G, top_n=3):
    """Analyze nodes with highest degree"""
    print("=== HIGH DEGREE NODES ANALYSIS ===\n")

    # Get degree for all nodes
    degrees = dict(G.degree())
    top_nodes = sorted(degrees.items(), key=lambda x: x[1], reverse=True)[:top_n]

    print(f"Top {top_n} nodes by degree:")
    for i, (node, degree) in enumerate(top_nodes):
        node_type = G.nodes[node].get("type", "unknown")
        print(f"  {i+1}. {node} (type: {node_type}, degree: {degree})")

        # Show some connections
        if G.is_directed():
            in_edges = list(G.in_edges(node))
            out_edges = list(G.out_edges(node))
            print(f"     In-degree: {len(in_edges)}, Out-degree: {len(out_edges)}")

            if in_edges:
                sample_in = [f"{u}->{node}" for u, _ in in_edges[:3]]
                print(f"     Sample incoming: {', '.join(sample_in)}")
            if out_edges:
                sample_out = [f"{node}->{v}" for _, v in out_edges[:3]]
                print(f"     Sample outgoing: {', '.join(sample_out)}")
        else:
            neighbors = list(G.neighbors(node))
            if neighbors:
                print(f"     Sample neighbors: {neighbors[:5]}")
        print()


def analyze_edge_attributes(G, max_examples=5):
    """Analyze edge attributes in detail"""
    print("=== EDGE ATTRIBUTES ANALYSIS ===\n")

    # Collect all edge attributes
    all_edge_attrs = defaultdict(list)

    if isinstance(G, (nx.MultiGraph, nx.MultiDiGraph)):
        for u, v, key, data in G.edges(keys=True, data=True):
            for attr_name, attr_value in data.items():
                all_edge_attrs[attr_name].append((u, v, key, attr_value))
    else:
        for u, v, data in G.edges(data=True):
            for attr_name, attr_value in data.items():
                all_edge_attrs[attr_name].append((u, v, attr_value))

    if not all_edge_attrs:
        print("No edge attributes found.")
        return

    print("Edge Attribute Summary:")
    for attr_name, values in all_edge_attrs.items():
        print(f"  {attr_name}: {len(values)} edges")

        # Show value distribution for this attribute
        if isinstance(values[0][-1], str):
            value_counts = Counter([item[-1] for item in values])
            print(f"    Value distribution: {dict(value_counts.most_common(5))}")
        print()

    # Show examples for each attribute
    for attr_name, values in all_edge_attrs.items():
        print(f"--- {attr_name.upper()} ATTRIBUTE EXAMPLES ---")

        for i, item in enumerate(values[:max_examples]):
            if isinstance(G, (nx.MultiGraph, nx.MultiDiGraph)):
                u, v, key, attr_value = item
                print(f"  {i+1}. {u} -> {v} (key: {key}): {attr_name} = {attr_value}")
            else:
                u, v, attr_value = item
                print(f"  {i+1}. {u} -> {v}: {attr_name} = {attr_value}")

        if len(values) > max_examples:
            print(f"     ... and {len(values) - max_examples} more")
        print()


def main():
    pickle_file = "/data_ext/ref_code/LocAgent/data/data/pp_data/pingp-auto-test.pkl"

    # Check if file exists
    if not os.path.exists(pickle_file):
        print(f"File does not exist: {pickle_file}")
        return

    print(f"Analyzing relationships in: {pickle_file}")
    print("=" * 80)

    # Load graph
    G = load_graph(pickle_file)
    if G is None:
        return

    # Basic info
    print(f"Graph type: {type(G)}")
    print(f"Nodes: {G.number_of_nodes()}, Edges: {G.number_of_edges()}")
    print(f"Directed: {G.is_directed()}, Multigraph: {G.is_multigraph()}")
    print()

    # Analyze relationships
    analyze_node_relationships(G, max_examples=5)

    # Analyze specific node types
    analyze_specific_node_types(G)

    # Analyze high degree nodes
    analyze_high_degree_nodes(G, top_n=15)

    # Analyze edge attributes
    analyze_edge_attributes(G, max_examples=10)


if __name__ == "__main__":
    main()
